
import lazyllm
from lazyllm import (
    fc_register, Document, Retriever,
    OnlineEmbeddingModule, OnlineChatModule, WebModule,
    ReactAgent
)

embed_model = OnlineEmbeddingModule(source="qwen")

doc = Document(dataset_path="D://a/demo/PythonProject2/dusheng/dsds", embed=embed_model)
retriever = Retriever(doc, group_name='CoarseChunk', similarity="bm25_chinese", topk=10)


# 注册RAG工具
@fc_register("tool")
def search_knowledge_base(query: str):
    """
    搜索知识库并返回相关文档内容

    Args:
        query (str): 搜索查询字符串
    """
    # 将Retriever组件召回的节点全部存储到列表doc_node_list中
    doc_node_list = retriever(query=query)
    # 将召回节点中的内容组合成字符串
    context_str = "".join([node.get_content() for node in doc_node_list])
    return context_str



# prompt 设计
prompt = '''
Continue writing the following content and use more unique metaphorical sentences, just like the previous writing style. To enrich the plot as much as possible and strengthen the character's personality traits. Give the character motivation and enrich the character image. (2) Narrative structure: The narrative is complete, and the event should include the complete cause, process, and result. The narrative rhythm is smooth. (3) Language style: Use straightforward and easy to understand words. Multiple segments, fewer long sentences, and strong emotions. (4) Concise white description. The expression method of direct narration in colloquial language does not require those "literary" or "artistic conception" things. The characters created do not need to be strictly described according to the real character's personality and behavior, and avoid extreme emotional atmosphere. Avoid using metaphors, object descriptions, and summary paragraphs, and instead focus more directly on the senses themselves, using more concrete words to present details. Avoid unnecessary rhetoric and psychological activities. Based on the provided materials or plot, pay attention to the authenticity of the story process, rather than the authenticity of specific data. Following the style of Shuangwen, use pure text
'''
# 创建ReactAgent)
agent = ReactAgent(
    llm=OnlineChatModule(source='qwen', model="qwen-plus", stream=False),
    tools=['search_knowledge_base'],
    prompt=prompt,
    stream=False
)

# 创建Web模块并启动
w = WebModule(agent, stream=False)
w.start().wait()
